#include <zlib.h>
#include "IPNGUtils.hxx"

#define ZLIB_BUF_SIZE 131072 // 128K

// PNG status codes returned by functions
#define PNG_SUCCESS			0L
#define PNG_FAIL			1L
#define PNG_INVALID_ARG			2L
#define PNG_WIN32_ERROR			3L
#define PNG_NOT_PNG			4L
#define PNG_INVALID_CRC			5L
#define PNG_BAD_FORMAT			6L
#define PNG_UNSUPPORTED			7L
#define PNG_NOT_ENOUGH_MEMORY		8L

// Filter sub types
#define PNG_FILTER_TYPE_NONE		0x00
#define PNG_FILTER_TYPE_SUB			0x01
#define PNG_FILTER_TYPE_UP			0x02
#define PNG_FILTER_TYPE_AVERAGE		0x03
#define PNG_FILTER_TYPE_PAETH		0x04

#define PNGMin(x,y) (x<=y?x:y)

// NOTE: GOOD articles for reference
// 0. A blog with good tables and summary: https://rimson.top/2019/08/23/png-format/
// 1. The png spec about file structure: http://www.libpng.org/pub/png/spec/1.2/PNG-Structure.html

unsigned long crc(unsigned char *buf, int len);

const uchar PNG_SIGNATURE[] = {137, 80, 78, 71, 13, 10, 26, 10};

// NOTE png needs big endian https://stackoverflow.com/questions/2384111/png-file-format-endianness
uint32_t to_uint32_big_endian(const uchar *a) {
    return a[3] + (((uint32_t) a[2]) << 8) + (((uint32_t) a[1]) << 16) + (((uint32_t) a[0]) << 24);
}

void from_uint32_big_endian(uchar *a, uint32_t n) {
    a[0] = (n >> 24) & 0xFF;
    a[1] = (n >> 16) & 0xFF;
    a[2] = (n >> 8) & 0xFF;
    a[3] = n & 0xFF;
}

int png_iterator_first_chunk(const vector<uchar> &bytes) {
    size_t n = sizeof(PNG_SIGNATURE) / sizeof(PNG_SIGNATURE[0]);
    for (int j = 0; j < n; ++j) {
        CV_Assert(bytes[j] == PNG_SIGNATURE[j]);
    }
    return n;
}

int png_iterator_next_chunk(const vector<uchar> &bytes, int chunk_begin_loc, string &chunk_type) {
    int idx = chunk_begin_loc;

    uint32_t length = to_uint32_big_endian(&bytes[idx]);
    idx += 4;

    string type(bytes.begin() + idx, bytes.begin() + idx + 4);
    idx += 4;

    chunk_type = type;
    return chunk_begin_loc + (int) length + 12;
}

vector<uchar> create_PLTE_chunk(const Mat &palette) {
    CV_Assert(palette.cols == 1);
    CV_Assert(palette.type() == CV_8UC3);
    CV_Assert(palette.rows <= 256); // since only 8-bit depth!

    // *3: since 3 bytes per channel (color image)
    uint32_t length = palette.rows * 3;

    // +12: by definition of chunk "actual" length vs "recorded" length
    vector<uchar> bytes(length + 12);

    int idx = 0;

    from_uint32_big_endian(&bytes[idx], length);
    idx += 4;

    bytes[idx + 0] = 'P';
    bytes[idx + 1] = 'L';
    bytes[idx + 2] = 'T';
    bytes[idx + 3] = 'E';
    idx += 4;

    // I know not that fast, but palette is pretty small so readability is also important
    for (int j = 0; j < palette.rows; ++j) {
        for (int ch = 0; ch < 3; ++ch) {
            bytes[idx + ch] = palette.at<Vec3b>(j, 0)[ch];
        }
        idx += 3;
    }

    uint32_t crc_val = crc(&bytes[4], (int) length + 4);
    from_uint32_big_endian(&bytes[idx], crc_val);
    idx += 4;

    CV_Assert(idx == bytes.size());

    return bytes;
}

void change_IHDR_colortype_and_crc(vector<uchar> &bytes, int ihdr_start_loc, int ihdr_end_loc) {
    const int ihdr_data_loc = ihdr_start_loc + 4 + 4;
    const int ihdr_bitdepth_loc = ihdr_data_loc + 8;
    const int ihdr_colortype_loc = ihdr_data_loc + 9;
    const int crc_loc = ihdr_end_loc - 4;

    // we need 8-bit depth
    CV_Assert(bytes[ihdr_bitdepth_loc] == 8);

    // a. change colortype to "indexed color"
    bytes[ihdr_colortype_loc] = 3;

    // b. re-calculate the CRC value
    int chunk_length = ihdr_end_loc - ihdr_start_loc - 12;
    uint32_t crc_val = crc(&bytes[ihdr_start_loc + 4], chunk_length + 4);
    from_uint32_big_endian(&bytes[crc_loc], crc_val);
}

// just return vec by value, it is preferred https://stackoverflow.com/questions/15704565/efficient-way-to-return-a-stdvector-in-c
// NOTE WARN: 注意bmp那边返回的是Mat而非vector<uchar>，好像是vector有个什么性能问题(?)
vector<uchar> imencode_palette_png(const Mat &im_index, const Mat &palette) {
    CV_Assert(im_index.type() == CV_8UC1);
    CV_Assert(palette.cols == 1);
    CV_Assert(palette.type() == CV_8UC3);

    vector<uchar> bytes;
    imencode(".png", im_index, bytes,
            // use higher compression, thus smaller
             {IMWRITE_PNG_COMPRESSION, 9});

    int idx = png_iterator_first_chunk(bytes);
    while (true) {
        string chunk_type;
        int next_chunk_start_idx = png_iterator_next_chunk(bytes, idx, chunk_type);

        if (chunk_type == "IHDR") {

            // change 1: change THDR "color" flag
            change_IHDR_colortype_and_crc(bytes, idx, next_chunk_start_idx);

            // change 2: insert the PLTE chunk **after** THDR chunk
            vector<uchar> plte_chunk_bytes = create_PLTE_chunk(palette);
            bytes.insert(bytes.begin() + next_chunk_start_idx, plte_chunk_bytes.begin(), plte_chunk_bytes.end());

            // no need to manipulate the data after that. let's stop!
            break;
        }

        if (next_chunk_start_idx >= bytes.size()) break;

        idx = next_chunk_start_idx;
    }

    return bytes;
}

// ---------------------------------------------------------------------------------------------------------
// NOTE copy from: **OFFICIAL** png documentation has a link to "sample crc code"
// https://www.w3.org/TR/PNG-CRCAppendix.html
// ---------------------------------------------------------------------------------------------------------

/* Table of CRCs of all 8-bit messages. */
unsigned long crc_table[256];

/* Flag: has the table been computed? Initially false. */
int crc_table_computed = 0;

/* Make the table for a fast CRC. */
void make_crc_table(void) {
    unsigned long c;
    int n, k;

    for (n = 0; n < 256; n++) {
        c = (unsigned long) n;
        for (k = 0; k < 8; k++) {
            if (c & 1)
                c = 0xedb88320L ^ (c >> 1);
            else
                c = c >> 1;
        }
        crc_table[n] = c;
    }
    crc_table_computed = 1;
}

/* Update a running CRC with the bytes buf[0..len-1]--the CRC
   should be initialized to all 1's, and the transmitted value
   is the 1's complement of the final running CRC (see the
   crc() routine below)). */

unsigned long update_crc(unsigned long crc, unsigned char *buf, int len) {
    unsigned long c = crc;
    int n;

    if (!crc_table_computed)
        make_crc_table();
    for (n = 0; n < len; n++) {
        c = crc_table[(c ^ buf[n]) & 0xff] ^ (c >> 8);
    }
    return c;
}

/* Return the CRC of the bytes buf[0..len-1]. */
unsigned long crc(unsigned char *buf, int len) {
    return update_crc(0xffffffffL, buf, len) ^ 0xffffffffL;
}

// ---------------------------------------------------------------------------------------------------------

/* 
 Uncompresses image data from IDAT
*/
long uncompress(
	unsigned char* pvIn,
	long dwSizeIn,
	unsigned char** ppvOut,
	long* pdwSizeOut
	)
{
	int ret;
	unsigned long have;
	z_stream strm;
	unsigned char* buf_in = NULL;
	unsigned char* buf_out = NULL;
	unsigned char* psrc = NULL;
	unsigned char* pdest = NULL;
	unsigned long srcSize = 0;
	unsigned long destSize = 0;
	unsigned long bytesRead = 0;
	unsigned long bytesWrite = 0;

	*ppvOut = NULL;
	*pdwSizeOut = 0;

	buf_in = (unsigned char*)malloc(ZLIB_BUF_SIZE);
	if (buf_in == NULL)
	{
		return PNG_NOT_ENOUGH_MEMORY;
	}

	buf_out = (unsigned char*)malloc(ZLIB_BUF_SIZE);
	if (buf_out == NULL)
	{
		free(buf_in);
		return PNG_NOT_ENOUGH_MEMORY;
	}

	// allocate memory for extraction destinate (final output)
	destSize = dwSizeIn;

	assert(destSize != 0);

	pdest = (unsigned char*)malloc(destSize);
	if (pdest == NULL)
	{
		free(buf_in);
		free(buf_out);
		return PNG_NOT_ENOUGH_MEMORY;
	}

	/* allocate inflate state */
	strm.zalloc = Z_NULL;
	strm.zfree = Z_NULL;
	strm.opaque = Z_NULL;
	strm.avail_in = 0;
	strm.next_in = Z_NULL;
	ret = inflateInit(&strm);
	if (ret != Z_OK)
	{
		free(buf_in);
		free(buf_out);
		free(pdest);
		return PNG_FAIL;
	}

	psrc = (unsigned char*)pvIn;
	srcSize = dwSizeIn;


	/* decompress until inflate stream ends */
	do {
		if (bytesRead < srcSize)
		{
			if (srcSize <= ZLIB_BUF_SIZE)
			{
				memcpy(buf_in, psrc, srcSize);
				strm.avail_in = srcSize;
				bytesRead = srcSize;
			}
			else
			{
				unsigned long n = PNGMin(ZLIB_BUF_SIZE, srcSize - bytesRead);
				memcpy(buf_in, psrc + bytesRead, n);
				bytesRead += n;
				strm.avail_in = n;
			}
		}


		if (strm.avail_in == 0)
			break;
		strm.next_in = buf_in;

		/* run inflate() on input until output buffer not full */
		do {
			strm.avail_out = ZLIB_BUF_SIZE;
			strm.next_out = buf_out;
			ret = inflate(&strm, Z_NO_FLUSH);

			assert(ret != Z_STREAM_ERROR);  /* state not clobbered */

			switch (ret) {
			case Z_NEED_DICT:
			case Z_DATA_ERROR:
			case Z_MEM_ERROR: {
					inflateEnd(&strm);
					free(buf_in);
					free(buf_out);
					free(pdest);
					return PNG_FAIL;
				}
			}

			have = ZLIB_BUF_SIZE - strm.avail_out;

			if (destSize >= bytesWrite + have)
			{
				memcpy(pdest + bytesWrite, buf_out, have);
				bytesWrite += have;
			}
			else // we need more memory
			{
				destSize += have;
				pdest = (unsigned char*)realloc(pdest, destSize);

				memcpy(pdest + bytesWrite, buf_out, have);
				bytesWrite += have;
			}
			
			/* done when inflate() says it's done */
		} while (strm.avail_out == 0);
		
	} while (ret != Z_STREAM_END);

	/* clean up */
	inflateEnd(&strm);
	free(buf_in);
	free(buf_out);

	if (ret == Z_STREAM_END)
	{
		// free extra memory
		if (bytesWrite < destSize)
		{
			destSize = bytesWrite;
			pdest = (unsigned char*)realloc(pdest, destSize);
		}

		*ppvOut = pdest;
		*pdwSizeOut = destSize;

		return PNG_SUCCESS;
	}

	free(pdest);

	return  PNG_FAIL;
}

/*
 Calculates and returns Paeth filter value
*/
int PaethPredictor(int a, int b, int c)
{
	int p = a + b - c;
	int pa = abs(p - a);
	int pb = abs(p - b);
	int pc = abs(p - c);

	int pr = 0;
	if (pa <= pb && pa <= pc)
		pr = a;
	else if (pb <= pc)
		pr = b;
	else
		pr = c;

	return pr;
}

long unfilterByte(
	uint lineNum, /*scanline number, 0 based*/ 
	uint bytePos, /*byte position in scanline*/
	uint type, /*filter type*/
	unsigned char* pbyte, /*filtered byte*/
	uint stride, /*stride of scanline*/
	uint bpp /*byte per pixel*/
	)
{
	char a = 0;
	char b = 0;
	char c = 0;

	switch (type)
	{
	case PNG_FILTER_TYPE_NONE:
		break;

	case PNG_FILTER_TYPE_SUB:
		if (bytePos >= bpp) {
			*pbyte += *(pbyte - bpp);
		}
		break;

	case PNG_FILTER_TYPE_UP:
		if (lineNum > 0) {
			*pbyte += *(pbyte - stride - 1);
		}
		break;

	case PNG_FILTER_TYPE_AVERAGE:
		if (bytePos >= bpp) {
			a = *(pbyte - bpp);
		}

		if (lineNum > 0) {
			b = *(pbyte - stride - 1);
		}

		*pbyte += (a + b) / 2;
		break;

	case PNG_FILTER_TYPE_PAETH:
		if (bytePos >= bpp) {
			a = *(pbyte - bpp);
		}

		if (lineNum > 0) {
			b = *(pbyte - stride - 1);
		}

		if (bytePos >= bpp && lineNum > 0) {
			c = *(pbyte - stride - 1 - bpp);
		}

		*pbyte += PaethPredictor(a, b, c);
		break;

	default:
		printf("UnfilterByte()- Wrong filter type, %d\n", type);
//		return PNG_BAD_FORMAT; // continue any way!
		break;
	}

	return PNG_SUCCESS;
}



bool imdecode_palette_png(char* fileData, size_t dataLength,  Mat &im_index, Mat& palette)
{
    vector<uchar> bytes(fileData, fileData + dataLength);

    png_iterator_first_chunk(bytes);

    int idx = png_iterator_first_chunk(bytes);

    int imageWidth = 0;
    int imageHeight = 0;

    cv::Mat indexMat;

    while (true) {
        string chunk_type;
        int next_chunk_start_idx = png_iterator_next_chunk(bytes, idx, chunk_type);
        
        printf("Next chunk start idx:%d\n", idx);

        if( chunk_type == "IDAT")
        {
            char buf[5];

            printf("Data Chunk\n");
            strncpy(buf,chunk_type.c_str(),4);
            buf[4]='\0';

            printf("Type:%s\n", buf);

            int val = to_uint32_big_endian(&bytes[idx]);
            printf("Data Length:%d\n", val);

            unsigned char* pv = 0;
            long pvs = 0;

            uncompress(reinterpret_cast<unsigned char *>(&bytes[idx + 8]), val, &pv, &pvs);
            printf("Uncompressed data length:%ld\n", pvs);

            int tWidth = imageWidth + 1;
            for(int i = 0; i < imageHeight; i ++)
            {               
               //printf("Row Filter Type: %d\t%d\n", i, pv[i * tWidth]);

               uint byteType = pv[i * tWidth];

               for( int j = 0; j < imageWidth; j ++)
               {
                    unfilterByte(i, j, byteType, &(pv[i * tWidth + j + 1]), imageWidth, 1 );

                    im_index.at<uchar>(i,j) = pv[ i * tWidth + j + 1];                    
               }
            }
        }
        else if( chunk_type == "PLTE")
        {
            char buf[5];

            printf("Platte Chunk\n");
            strncpy(buf,chunk_type.c_str(),4);
            buf[4]='\0';

            printf("Type:%s\n", buf);

            int val = to_uint32_big_endian(&bytes[idx]);
            printf("Data Length:%d\n", val);

            int colorNumber = val / 3;

            palette = cv::Mat(colorNumber,1,CV_8UC3,cv::Scalar(0,0,0));

            for(int i = 0; i < colorNumber; i ++)
            {
                Vec3b color = cv::Vec3b(0,0,0);

                color[0] = bytes[idx + 8 + i * 3 ];
                color[1] = bytes[idx + 8 + i * 3 + 1];
                color[2] = bytes[idx + 8 + i * 3 + 2];

                palette.at<cv::Vec3b>(i,0) = color;

                printf("Color %d:\t(%d,%d,%d)\n",i, color[0], color[1], color[2]);

            }
        }
        else if( chunk_type == "IHDR" )
        {
            printf("-------------------Header Chunk-------------------\n");

            int val = to_uint32_big_endian(&bytes[idx]);
            printf("Data Length:%d\n", val);

            imageWidth = to_uint32_big_endian(&bytes[idx + 8]);
            imageHeight = to_uint32_big_endian(&bytes[idx + 12]);

            im_index = cv::Mat(cvSize(imageWidth,imageHeight), CV_8UC1, cv::Scalar(0,0,0));

            uchar bitDepth = bytes[idx + 16];
            uchar colorType = bytes[idx + 17];
            uchar compressionMethod = bytes[idx + 18];
            uchar filterMethod = bytes[idx + 19];
            uchar interlaceMethod = bytes[idx + 20];

            printf("Header Width:%d\n", imageWidth);
            printf("Header Height:%d\n",imageHeight);
            printf("Bit Depth:%d\n", bitDepth);
            printf("Color Type:%d\n", colorType);
            printf("Compression Method:%d\n", compressionMethod);
            printf("Filter Method:%d\n", compressionMethod);
            printf("Interlace Method:%d\n", interlaceMethod);
        }
        else
        {
            char buf[5];

            printf("-------------------Other Chunk-------------------\n");
            strncpy(buf,chunk_type.c_str(),4);
            buf[4]='\0';

            printf("Type:%s\n", buf);

            printf("Other Chunk\n");
        }

        if (next_chunk_start_idx >= bytes.size()) break;

        idx = next_chunk_start_idx;
    }
}